import os
import shutil
import yaml
from pathlib import Path
from typing import Dict, Any

class BatchFileFilterProcessor:
    def __init__(self, input_folder: str, output_folder: str, config_file: str = None):
        self.input_folder = Path(input_folder)
        self.output_folder = Path(output_folder)
        self.config_file = Path(config_file) if config_file else None
        self.config = {}
        self.errors = []

    def load_config(self):
        if self.config_file is None or not self.config_file.exists():
            # 使用默认配置文件
            default_config = Path(__file__).parent / 'configs' / 'batch_file_filter_default_text.yml'
            if not default_config.exists():
                raise FileNotFoundError(f"Default config file not found: {default_config}")
            self.config_file = default_config
            print(f"使用默认的配置文件 Using default config file: {self.config_file}")

        with open(self.config_file, 'r') as file:
            self.config = yaml.safe_load(file)

    def process(self):
        # 检查输入文件夹是否存在
        if not self.input_folder.exists():
            raise FileNotFoundError(f"Input folder not found: {self.input_folder}")

        # 确保输出文件夹存在
        self.output_folder.mkdir(parents=True, exist_ok=True)

        # 加载配置文件
        self.load_config()

        # 遍历输入文件夹
        for root, _, files in os.walk(self.input_folder):
            for file in files:
                input_file = Path(root) / file
                relative_path = input_file.relative_to(self.input_folder)
                output_file = self.output_folder / relative_path

                try:
                    if self._filter_file(input_file):
                        self._copy_file(input_file, output_file)
                except Exception as e:
                    self.errors.append(f"Error processing {input_file}: {str(e)}")

        # 输出所有异常
        if self.errors:
            print("Errors encountered during processing:")
            for error in self.errors:
                print(error)

    def _filter_file(self, file: Path) -> bool:
        # 检查文件大小
        if 'size' in self.config:
            size_limit = self._parse_size(self.config['size']['less_equal'])
            if file.stat().st_size > size_limit:
                return False

        # 检查文件名
        if 'filename' in self.config:
            if 'not_include' in self.config['filename']:
                for keyword in self.config['filename']['not_include']:
                    if keyword in file.name:
                        return False

        # 检查文件类型
        if 'filetype' in self.config:
            if 'include' in self.config['filetype']:
                if file.suffix[1:] not in self.config['filetype']['include']:
                    return False

        return True

    def _parse_size(self, size_str: str) -> int:
        units = {'K': 1024, 'M': 1024**2, 'G': 1024**3}
        number = int(size_str[:-1])
        unit = size_str[-1].upper()
        return number * units.get(unit, 1)

    def _copy_file(self, src: Path, dst: Path):
        dst.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(src, dst)
        print(f"Copied: {src} -> {dst}")

def main():
    import argparse

    parser = argparse.ArgumentParser(description="Batch filter and copy files based on config.")
    parser.add_argument("-i", "--input", required=True, help="Input folder containing documents")
    parser.add_argument("-o", "--output", required=True, help="Output folder for filtered files")
    parser.add_argument("-c", "--config", help="Path to the YAML config file (optional)")
    args = parser.parse_args()

    processor = BatchFileFilterProcessor(args.input, args.output, args.config)
    processor.process()

if __name__ == "__main__":
    main()

"""
使用示例：
# 使用默认配置文件
python3 -m document2llm.io.batch_file_filter -i input_documents -o output_filtered

# 使用自定义配置文件
python3 -m document2llm.io.batch_file_filter -i input_documents -o output_filtered -c custom_config.yml

或者在Python代码中：

from document2llm.io.batch_file_filter import BatchFileFilterProcessor

# 使用默认配置文件
processor = BatchFileFilterProcessor("input_documents", "output_filtered")
processor.process()

# 使用自定义配置文件
processor = BatchFileFilterProcessor("input_documents", "output_filtered", "custom_config.yml")
processor.process()
"""






#
